package main

import (
	"context"
	"flag"
	"google.golang.org/grpc"
	"log"
	"pfs/profiling"
)

var ip, port string
var job string // namespace/jobName, such as default/my-job-name

func main() {
	flag.StringVar(&ip, "ip", "127.0.0.1", "ClusterD srv ip")
	flag.StringVar(&port, "port", "8899", "ClusterD srv port")
	flag.StringVar(&job, "job", "default", "job name space")
	flag.Parse()

	log.Printf("job name with namespace=%s\n", job)
	address := ip + ":" + port
	log.Printf("client will connect to %s\n", address)

	conn, err := grpc.Dial(address, grpc.WithInsecure())
	if err != nil {
		log.Fatalf("connect to %s failed, err=%v", address, err)
	}
	defer func() {
		closeErr := conn.Close()
		if closeErr != nil {
			log.Println(closeErr.Error())
		}
	}()
	client := profiling.NewTrainingDataTraceClient(conn)
	modifyData := &profiling.DataTypeReq{
		JobNsName: job,
		ProfilingSwitch: &profiling.ProfilingSwitch{
			CommunicationOperator: "off",
			Step:                  "off",
			SaveCheckpoint:        "off",
			FP:                    "off",
			DataLoader:            "off",
		},
	}
	_, err = client.ModifyTrainingDataTraceSwitch(context.Background(), modifyData)
	if err != nil {
		log.Fatalf("failed to change profiling switch, err=%v", err)
	}
	log.Printf("success change profiling switch to %v", modifyData.ProfilingSwitch)
}
